import os
import numpy as np
import cv2
import re


def extract_prefix(filename):
    pattern = re.compile(r'(.*)tr\d+-tc\d+\.png')
    match = pattern.match(filename)
    if match:
        return match.group(1)
    return ""


def cut_and_save_large_image(image_dir, output_dir):
    os.makedirs(output_dir, exist_ok=True)

    image_files = sorted([f for f in os.listdir(image_dir) if f.endswith('.png')])

    for image_file in image_files:
        prefix = extract_prefix(image_file)
        # 解析行和列编号
        match = re.search(r'tr(\d+)-tc(\d+)', image_file)
        if not match:
            continue

        big_tr = int(match.group(1))
        big_tc = int(match.group(2))

        image_path = os.path.join(image_dir, image_file)
        large_image = cv2.imread(image_path)
        if large_image is None:
            raise FileNotFoundError(f"Unable to open image file: {image_path}")

        large_tile_size = 1024

        # 计算大图的起始和结束小图编号
        if big_tr == 1:
            row_start, row_end = 1, 15
            row_valid_start, row_valid_end = 1, 13
        elif big_tr == 2:
            row_start, row_end = 12, 26
            row_valid_start, row_valid_end = 14, 24
        else:
            row_start, row_end = 23, 37
            row_valid_start, row_valid_end = 25, 37

        if big_tc == 1:
            col_start, col_end = 1, 13
            col_valid_start, col_valid_end = 1, 12
        elif big_tc == 2:
            col_start, col_end = 11, 23
            col_valid_start, col_valid_end = 13, 21
        else:
            col_start, col_end = 21, 33
            col_valid_start, col_valid_end = 22, 33

        for small_row in range(row_start, row_end + 1):
            if small_row < row_valid_start or small_row > row_valid_end:
                continue
            for small_col in range(col_start, col_end + 1):
                if small_col < col_valid_start or small_col > col_valid_end:
                    continue

                x_offset = (small_row - row_start) * large_tile_size
                y_offset = (small_col - col_start) * large_tile_size

                small_image = large_image[x_offset:x_offset + large_tile_size, y_offset:y_offset + large_tile_size]

                small_filename = f"{prefix}tr{small_row}-tc{small_col}.png"
                small_image_path = os.path.join(output_dir, small_filename)
                cv2.imwrite(small_image_path, small_image)
                print(f"Saved {small_filename} to {small_image_path}")

# 示例用法
image_dir = 'E:\\wafer52\\11867_32nm_big_fine_align2(1)'  # 替换为实际的大图文件夹路径
output_dir = 'E:\\wafer52\\11867_32nm_fine_align2(1)'  # 替换为实际的小图输出文件夹路径

cut_and_save_large_image(image_dir, output_dir)
